'''
@File    :   tree_plotter.py
@Version : 1.0
@Author :   iherr
@Desciption : 绘制决策树
'''

import matplotlib.pyplot as plt
import shang
import trees

decisionNode = dict(boxstyle="sawtooth", fc="0.8")
leafNode = dict(boxstyle="round4", fc="0.8")
arrow_args = dict(arrowstyle="<-")


def getNumLeafs(myTree):
    '''
    获取决策树的叶子结点的数目
    :param myTree:决策树
    :return: 决策树的叶子结点的数目
    '''
    numLeafs = 0
    firstStr = next(iter(myTree))  #list(myTree.keys())[0]
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[
                    key]).__name__ == 'dict': #检测是否为叶子节点
            numLeafs += getNumLeafs(secondDict[key])
        else:
            numLeafs += 1
    return numLeafs


def getTreeDepth(myTree):
    '''
    获取决策树的层数（深度）
    :param myTree:决策树
    :return:决策树的层数（深度）
    '''
    maxDepth = 0
    firstStr = next(iter(myTree))
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[
                    key]).__name__ == 'dict':  #检测是否为叶子节点
            thisDepth = 1 + getTreeDepth(secondDict[key])
        else:
            thisDepth = 1
        if thisDepth > maxDepth: maxDepth = thisDepth
    return maxDepth


def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    '''
    绘制结点
    :param nodeTxt:节点名称
    :param centerPt:文本位置
    :param parentPt:标注的箭头位置
    :param nodeType:节点格式
    :return:
    '''
    createPlot.ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',
                            xytext=centerPt, textcoords='axes fraction',
                            va="center", ha="center", bbox=nodeType, arrowprops=arrow_args)


def plotMidText(cntrPt, parentPt, txtString):
    '''
    标注有向边属性值
    :param cntrPt: 标注位置
    :param parentPt:标注位置
    :param txtString: 标注的内容
    :return:
    '''
    xMid = (parentPt[0] - cntrPt[0]) / 2.0 + cntrPt[0]
    yMid = (parentPt[1] - cntrPt[1]) / 2.0 + cntrPt[1]
    createPlot.ax1.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)


def plotTree(myTree, parentPt, nodeTxt):
    numLeafs = getNumLeafs(myTree)  # 获取叶子节点数
    depth = getTreeDepth(myTree)    # 获取层数
    firstStr = next(iter(myTree))  # 节点的文本内容
    cntrPt = (plotTree.xOff + (1.0 + float(numLeafs)) / 2.0 / plotTree.totalW, plotTree.yOff) #中心位置
    plotMidText(cntrPt, parentPt, nodeTxt) #标注有向边属性
    plotNode(firstStr, cntrPt, parentPt, decisionNode) #绘制节点
    secondDict = myTree[firstStr]
    plotTree.yOff = plotTree.yOff - 1.0 / plotTree.totalD
    for key in secondDict.keys():
        if type(secondDict[
                    key]).__name__ == 'dict':  # 如果不是叶子节点
            plotTree(secondDict[key], cntrPt, str(key))  # 递归调用
        else:  # 绘制叶子节点
            plotTree.xOff = plotTree.xOff + 1.0 / plotTree.totalW
            plotNode(secondDict[key], (plotTree.xOff, plotTree.yOff), cntrPt, leafNode)
            plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(key))
    plotTree.yOff = plotTree.yOff + 1.0 / plotTree.totalD


def createPlot(inTree):
    fig = plt.figure(1, facecolor='white') #创建
    fig.clf()   #清空
    axprops = dict(xticks=[], yticks=[])
    createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)  # 无xy轴
    plotTree.totalW = float(getNumLeafs(inTree))
    plotTree.totalD = float(getTreeDepth(inTree))
    plotTree.xOff = -0.5 / plotTree.totalW;
    plotTree.yOff = 1.0;
    plotTree(inTree, (0.5, 1.0), '') #绘制决策树
    plt.show()  #显示绘图

if __name__ == '__main__':
    dataSet,label=shang.createDataSet()
    myTree=trees.createTree(dataSet,label)
    createPlot(myTree)